"""
Wan 2.2 API Client

This module provides a Python client for interacting with the Wan 2.2 Video Generation API.
"""
import os
import time
import httpx
from typing import Dict, Optional, Union, List, Any
from enum import Enum
from pydantic import BaseModel, Field

class TaskStatus(str, Enum):
    QUEUED = "queued"
    PROCESSING = "processing"
    COMPLETED = "completed"
    FAILED = "failed"

class TaskType(str, Enum):
    T2V = "T2V"  # Text to Video
    I2V = "I2V"  # Image to Video
    TI2V = "TI2V"  # Text + Image to Video

class BaseTaskRequest(BaseModel):
    """Base model for task requests."""
    task_type: TaskType
    duration_seconds: int = Field(..., gt=0, le=60, description="Duration of the video in seconds (1-60)")
    resolution: str = Field("720p", regex=r'^\d{3,4}[pP]$', description="Output resolution (e.g., 720p, 1080p)")
    fps: int = Field(24, ge=1, le=60, description="Frames per second (1-60)")
    seed: Optional[int] = Field(None, description="Random seed for reproducibility")
    callback_url: Optional[str] = Field(None, description="URL to receive task completion webhook")

class T2VRequest(BaseTaskRequest):
    """Text-to-Video task request."""
    task_type: TaskType = TaskType.T2V
    prompt: str = Field(..., min_length=1, max_length=1000, description="Text prompt for video generation")
    negative_prompt: Optional[str] = Field(None, description="Text to avoid in the generated video")

class I2VRequest(BaseTaskRequest):
    """Image-to-Video task request."""
    task_type: TaskType = TaskType.I2V
    image_url: str = Field(..., description="URL of the input image")
    motion_strength: float = Field(0.5, ge=0.0, le=1.0, description="Strength of motion effect (0.0-1.0)")

class TI2VRequest(BaseTaskRequest):
    """Text+Image-to-Video task request."""
    task_type: TaskType = TaskType.TI2V
    prompt: str = Field(..., min_length=1, max_length=1000, description="Text prompt for video generation")
    image_url: str = Field(..., description="URL of the input image")
    negative_prompt: Optional[str] = Field(None, description="Text to avoid in the generated video")
    motion_strength: float = Field(0.5, ge=0.0, le=1.0, description="Strength of motion effect (0.0-1.0)")

class TaskResponse(BaseModel):
    """Response model for task operations."""
    task_id: str
    status: TaskStatus
    created_at: float
    updated_at: float
    result: Optional[Dict[str, Any]] = None
    error: Optional[str] = None

class WanClient:
    """Client for interacting with the Wan 2.2 Video Generation API."""
    
    def __init__(self, base_url: str = "http://localhost:8000", api_key: Optional[str] = None):
        """Initialize the Wan API client.
        
        Args:
            base_url: Base URL of the Wan API server
            api_key: Optional API key for authentication
        """
        self.base_url = base_url.rstrip('/')
        self.api_key = api_key or os.getenv('WAN_API_KEY')
        self.headers = {"Content-Type": "application/json"}
        if self.api_key:
            self.headers["Authorization"] = f"Bearer {self.api_key}"
    
    async def _request(self, method: str, endpoint: str, **kwargs) -> Dict[str, Any]:
        """Make an HTTP request to the API."""
        url = f"{self.base_url}/{endpoint.lstrip('/')}"
        async with httpx.AsyncClient() as client:
            response = await client.request(
                method,
                url,
                headers=self.headers,
                **kwargs
            )
            response.raise_for_status()
            return response.json()
    
    async def submit_task(self, task: Union[T2VRequest, I2VRequest, TI2VRequest]) -> TaskResponse:
        """Submit a new video generation task.
        
        Args:
            task: Task request object (T2VRequest, I2VRequest, or TI2VRequest)
            
        Returns:
            TaskResponse with task details
        """
        data = await self._request("POST", "/api/tasks", json=task.dict())
        return TaskResponse(**data)
    
    async def get_task_status(self, task_id: str) -> TaskResponse:
        """Get the status of a task.
        
        Args:
            task_id: ID of the task to check
            
        Returns:
            TaskResponse with current status and result if available
        """
        data = await self._request("GET", f"/api/tasks/{task_id}")
        return TaskResponse(**data)
    
    async def wait_for_task_completion(
        self, 
        task_id: str, 
        poll_interval: float = 5.0,
        timeout: float = 3600.0
    ) -> TaskResponse:
        """Wait for a task to complete.
        
        Args:
            task_id: ID of the task to wait for
            poll_interval: How often to poll for updates (seconds)
            timeout: Maximum time to wait (seconds)
            
        Returns:
            Final TaskResponse when task completes or times out
            
        Raises:
            TimeoutError: If the task does not complete within the timeout
        """
        start_time = time.time()
        
        while True:
            # Check timeout
            elapsed = time.time() - start_time
            if elapsed > timeout:
                raise TimeoutError(f"Task {task_id} did not complete within {timeout} seconds")
            
            # Get current status
            status = await self.get_task_status(task_id)
            
            # Return if task is complete or failed
            if status.status in (TaskStatus.COMPLETED, TaskStatus.FAILED):
                return status
            
            # Wait before polling again
            await asyncio.sleep(poll_interval)
    
    # Convenience methods for each task type
    
    async def text_to_video(
        self,
        prompt: str,
        duration_seconds: int = 5,
        resolution: str = "720p",
        fps: int = 24,
        negative_prompt: Optional[str] = None,
        seed: Optional[int] = None,
        callback_url: Optional[str] = None,
        wait: bool = False,
        **kwargs
    ) -> Union[TaskResponse, Dict[str, Any]]:
        """Generate a video from text.
        
        Args:
            prompt: Text prompt for video generation
            duration_seconds: Duration of the video in seconds (1-60)
            resolution: Output resolution (e.g., 720p, 1080p)
            fps: Frames per second (1-60)
            negative_prompt: Text to avoid in the generated video
            seed: Random seed for reproducibility
            callback_url: URL to receive task completion webhook
            wait: If True, wait for task completion and return the result
            **kwargs: Additional arguments to pass to the API
            
        Returns:
            TaskResponse or the final result if wait=True
        """
        task = T2VRequest(
            prompt=prompt,
            duration_seconds=duration_seconds,
            resolution=resolution,
            fps=fps,
            negative_prompt=negative_prompt,
            seed=seed,
            callback_url=callback_url,
            **kwargs
        )
        
        response = await self.submit_task(task)
        
        if wait:
            response = await self.wait_for_task_completion(response.task_id)
            if response.status == TaskStatus.COMPLETED:
                return response.result
            else:
                raise Exception(f"Task failed: {response.error}")
        
        return response
    
    async def image_to_video(
        self,
        image_url: str,
        duration_seconds: int = 5,
        resolution: str = "720p",
        fps: int = 24,
        motion_strength: float = 0.5,
        seed: Optional[int] = None,
        callback_url: Optional[str] = None,
        wait: bool = False,
        **kwargs
    ) -> Union[TaskResponse, Dict[str, Any]]:
        """Animate an image into a video.
        
        Args:
            image_url: URL of the input image
            duration_seconds: Duration of the video in seconds (1-60)
            resolution: Output resolution (e.g., 720p, 1080p)
            fps: Frames per second (1-60)
            motion_strength: Strength of motion effect (0.0-1.0)
            seed: Random seed for reproducibility
            callback_url: URL to receive task completion webhook
            wait: If True, wait for task completion and return the result
            **kwargs: Additional arguments to pass to the API
            
        Returns:
            TaskResponse or the final result if wait=True
        """
        task = I2VRequest(
            image_url=image_url,
            duration_seconds=duration_seconds,
            resolution=resolution,
            fps=fps,
            motion_strength=motion_strength,
            seed=seed,
            callback_url=callback_url,
            **kwargs
        )
        
        response = await self.submit_task(task)
        
        if wait:
            response = await self.wait_for_task_completion(response.task_id)
            if response.status == TaskStatus.COMPLETED:
                return response.result
            else:
                raise Exception(f"Task failed: {response.error}")
        
        return response
    
    async def text_and_image_to_video(
        self,
        prompt: str,
        image_url: str,
        duration_seconds: int = 5,
        resolution: str = "720p",
        fps: int = 24,
        motion_strength: float = 0.5,
        negative_prompt: Optional[str] = None,
        seed: Optional[int] = None,
        callback_url: Optional[str] = None,
        wait: bool = False,
        **kwargs
    ) -> Union[TaskResponse, Dict[str, Any]]:
        """Generate a video from text and an image.
        
        Args:
            prompt: Text prompt for video generation
            image_url: URL of the input image
            duration_seconds: Duration of the video in seconds (1-60)
            resolution: Output resolution (e.g., 720p, 1080p)
            fps: Frames per second (1-60)
            motion_strength: Strength of motion effect (0.0-1.0)
            negative_prompt: Text to avoid in the generated video
            seed: Random seed for reproducibility
            callback_url: URL to receive task completion webhook
            wait: If True, wait for task completion and return the result
            **kwargs: Additional arguments to pass to the API
            
        Returns:
            TaskResponse or the final result if wait=True
        """
        task = TI2VRequest(
            prompt=prompt,
            image_url=image_url,
            duration_seconds=duration_seconds,
            resolution=resolution,
            fps=fps,
            motion_strength=motion_strength,
            negative_prompt=negative_prompt,
            seed=seed,
            callback_url=callback_url,
            **kwargs
        )
        
        response = await self.submit_task(task)
        
        if wait:
            response = await self.wait_for_task_completion(response.task_id)
            if response.status == TaskStatus.COMPLETED:
                return response.result
            else:
                raise Exception(f"Task failed: {response.error}")
        
        return response

# Example usage
async def example():
    # Initialize client
    client = WanClient(base_url="http://localhost:8000")
    
    # Example 1: Text to Video
    print("Submitting text-to-video task...")
    task = await client.text_to_video(
        prompt="A beautiful sunset over mountains, cinematic, 8K",
        duration_seconds=5,
        resolution="720p"
    )
    print(f"Task submitted: {task.task_id}")
    
    # Wait for completion
    try:
        result = await client.wait_for_task_completion(task.task_id)
        print(f"Task completed! Result: {result}")
    except Exception as e:
        print(f"Error: {e}")

if __name__ == "__main__":
    import asyncio
    asyncio.run(example())
